#ifndef __GRABCUT_
#define __GRABCUT_

#include "GMM.hpp"

#include "GCGraph.hpp"
#include <limits>

class GrabCut
{
private:
    double calcBeta(const cv::Mat &img);
    void calcNWeights(const cv::Mat &img, cv::Mat &leftW, cv::Mat &upleftW, cv::Mat &upW,
        cv::Mat &uprightW, double beta, double gamma);
    void checkMask(const cv::Mat &img, const cv::Mat &mask);
    void initMaskWithRect(cv::Mat &mask, cv::Size imgSize, cv::Rect rect);
    void initGMMs(const cv::Mat &img, const cv::Mat &mask, GMM &bgdGMM, GMM &fgdGMM);
    void assignGMMsComponents(const cv::Mat &img, const cv::Mat &mask, const GMM &bgdGMM,
                              const GMM &fgdGMM, cv::Mat &compIdxs);
    void learnGMMs(const cv::Mat &img, const cv::Mat &mask, const cv::Mat &compIdxs, GMM &bgdGMM, GMM &fgdGMM);
    void constructGCGraph(const cv::Mat &img, const cv::Mat &mask, const GMM &bgdGMM, const GMM &fgdGMM,
                          double lambda, const cv::Mat &leftW, const cv::Mat &upleftW, const cv::Mat &upW,
                          const cv::Mat &uprightW, GCGraph &graph);
    void estimateSegmentation(GCGraph &graph, cv::Mat &mask);

public:
    GrabCut(){};
    ~GrabCut(){};

    void operator()(cv::Mat _img, cv::Mat _mask, cv::Rect rect,
        cv::Mat _bgdModel, cv::Mat _fgdModel,
                    int iterCount, int mode);
};

void changeBackGD(cv::Mat& forGD, cv::Mat& backGD, cv::Mat& mask, cv::Mat& dst, bool flag);

#endif